{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import scared"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Leakage Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import scared\n",
    "@scared.attack_selection_function\n",
    "def key_pt_xor(plaintext, guesses):\n",
    "    res = np.empty((plaintext.shape[0], len(guesses), plaintext.shape[1]), dtype='uint8')\n",
    "    for i, guess in enumerate(guesses):\n",
    "        res[:,i,:] = np.bitwise_xor(plaintext, guess)\n",
    "    return res\n",
    "\n",
    "@scared.attack_selection_function\n",
    "def sbox_hd(plaintext, guesses):\n",
    "    res = np.empty((plaintext.shape[0], len(guesses), plaintext.shape[1]), dtype='uint8')\n",
    "    for i, guess in enumerate(guesses):\n",
    "        res[:,i,:] = np.bitwise_xor(scared.aes.sub_bytes(np.bitwise_xor(plaintext, guess)), np.bitwise_xor(plaintext, guess))\n",
    "    return res\n",
    "\n",
    "@scared.attack_selection_function\n",
    "def sbox_hw(plaintext, guesses):\n",
    "    res = np.empty((plaintext.shape[0], len(guesses), plaintext.shape[1]), dtype='uint8')\n",
    "    for i, guess in enumerate(guesses):\n",
    "        res[:,i,:] = scared.aes.sub_bytes(np.bitwise_xor(plaintext, guess))\n",
    "    return res\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## T_Table"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scared.models import HammingWeight\n",
    "@numba.vectorize([numba.uint32(numba.uint8)])\n",
    "def t_table_vec(x):\n",
    "    return _T_TABLE_LUT[x]\n",
    "\n",
    "class TTable(HammingWeight):\n",
    "    def _compute(self, data, axis):\n",
    "        if data.dtype.kind != 'u':\n",
    "            raise ValueError(f'HammingWeight should take unsigned integer data as input, not {data.dtype}).')\n",
    "\n",
    "        if data.dtype != self.expected_dtype:\n",
    "            raise ValueError(f'Expected dtype for HammingWeight input data is {self.expected_dtype}, not {data.dtype}.')\n",
    "\n",
    "        if data.shape[axis] < self.nb_words:\n",
    "            raise ValueError(f'data should have at least {self.nb_words} as dimension with index {axis}, not {data.shape[axis]}.')\n",
    "\n",
    "        result_data = t_table_vec(data)\n",
    "        if self.nb_words > 1:\n",
    "            final_w_dimension = data.shape[axis] // self.nb_words\n",
    "            final_shape = [d if i != axis else final_w_dimension for i, d in enumerate(data.shape)]\n",
    "            result = _np.zeros(final_shape, dtype='uint32').swapaxes(0, axis)\n",
    "            result_data = _np.swapaxes(result_data, 0, axis)\n",
    "            for i in range(result.shape[0]):\n",
    "                slices = result_data[i * self.nb_words: (i + 1) * self.nb_words]\n",
    "                result[i] = _np.sum(slices, axis=0)\n",
    "            result_data = result\n",
    "            result_data = _np.swapaxes(result, 0, axis)\n",
    "\n",
    "        return result_data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## CW Project to estraces"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import estraces\n",
    "#del cw_traces\n",
    "cw_traces = estraces.read_ths_from_ram(np.array(project.waves), \n",
    "                                      plaintext=np.array(project.textins),\n",
    "                                      ciphertext=np.array(project.textouts))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Attack Object"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "container = scared.Container(cw_traces)\n",
    "a = scared.CPAAttack(selection_function=sbox_hd,\n",
    "                    model=HammingWeight(),\n",
    "                    discriminant=scared.maxabs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Running Attack and Printing Results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "a.run(container)\n",
    "for byte in np.argmax(a.scores, axis=0):\n",
    "    print(\"{:02X}\".format(byte%0xFF))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Print Results and Plot Correlation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from bokeh.plotting import figure, show\n",
    "from bokeh.io import output_notebook\n",
    "output_notebook()\n",
    "p = figure()\n",
    "key_guess = []\n",
    "for i in range(16):\n",
    "    results = a.results\n",
    "    print(\"Correct key: {:02X} (corr= {})\".format(project.keys[0][i], a.scores[project.keys[0][i]][i]))\n",
    "    sorted_scores = np.argsort(a.scores, axis=0)[::-1]\n",
    "    j=0\n",
    "    pge = 0\n",
    "    for score in sorted_scores:\n",
    "        #print(hex(score[i]), a.scores[score[i]][i])\n",
    "        if score[i] == project.keys[0][i]:\n",
    "            pge = j \n",
    "        j += 1\n",
    "    print(\"PGE={}\".format(pge))\n",
    "    guess = sorted_scores[0][i]\n",
    "    sec_guess = sorted_scores[1][i]\n",
    "    third_guess = sorted_scores[2][i]\n",
    "    xrange = list(range(len(results[guess][i])))\n",
    "    print(\"Best Guess is {:02X} ({:02X}) (Corr = {})\".format(guess, project.keys[0][i], a.scores[guess][i]))\n",
    "    #print(\"Second Guess is {:02X} ({:02X}) (Corr = {})\".format(sec_guess, project.keys[0][i], a.scores[sec_guess][i]))\n",
    "    #print(\"Third Guess is {:02X} ({:02X}) (Corr = {})\".format(third_guess, project.keys[0][i], a.scores[third_guess][i]))\n",
    "    p.line(xrange, results[guess][i], color=\"red\")\n",
    "    #p.line(xrange, results[sec_guess][i], color=\"green\")\n",
    "    #p.line(xrange, results[third_guess][i], color=\"green\")\n",
    "    p.line(xrange, results[project.keys[0][i]][i], color=\"green\")\n",
    "    key_guess.append(guess)\n",
    "    \n",
    "show(p)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
